package rerac.feature;

import java.util.ArrayList;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.Map.Entry;

import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;

import rerac.protos.InstanceCollection.FeatureMap;
import rerac.protos.InstanceCollection.Instance;
import rerac.protos.InstanceCollection.Instance.Builder;
import rerac.protos.InstanceCollection.Instance.Feature;
import rerac.protos.Corpus.Document;
import rerac.protos.Corpus.Document.Compound;
import rerac.protos.Corpus.Document.CompoundGroup;
import rerac.protos.Corpus.Document.AnnotationType;

/**
 * Class for extracting feature counts from documents.
 * Features are strings, but they can be stored as hashes as well in order to 
 * save memory. In this case it may happen (although with a small chance only) 
 * that two different features are indistinguishable as they have the same hash.
 * 
 * See extract() for a description of the features extracted.
 * 
 * @author Benjamin Roth
 *
 */
// TODO: Include per seed-pair training examples, rather than per extraction.
public class SentenceFeatureExtractor {
  BiMap<String, Integer> features = HashBiMap.create();
  BiMap<Integer, Integer> featureHashes = HashBiMap.create();
  BiMap<String, Integer> labels = HashBiMap.create();
  boolean useHash = false;

  /**
   * 
   * @param useFeatureHash whether features are represented as full strings 
   * (true) or hashes (false).
   */
  public SentenceFeatureExtractor(boolean useFeatureHash) {
    useHash = useFeatureHash;
  }

  /**
   * This sets the class label of the instance for which features are to be 
   * extracted.
   * 
   * @param instance
   * @param label
   */
  private void setLabel(Builder instance, String label) {
    Integer labelIdx = labels.size();
    if (labels.containsKey(label)) {
      labelIdx = labels.get(label);
    } else {
      labels.put(label, labelIdx);
    }
    instance.setCategoryIndex(labelIdx);
    instance.setCategory(label);
    instance.setPositive(true);
  }
  
  /**
   * Returns the index of the feature; a new index is added if the feature 
   * occurs the first time.
   * This methods is to be used if useHash == false.
   * 
   * @param fStr feature string
   * @return the index
   */
  private int getFeatureIdx(String fStr) {
    Integer featureIdx = features.size();
    if (features.containsKey(fStr)) {
      featureIdx = features.get(fStr);
    } else {
      features.put(fStr, featureIdx);
    }
    return featureIdx;
  }

  /**
   * Returns the index of hash of a feature; a new index is added if the 
   * corresponding hash occurs the first time.
   * This methods is to be used if useHash == true.
   * 
   * @param fStr feature string
   * @return the index
   */
  private int getFeatureHashIdx(String fStr) {
    int fHash = fStr.hashCode();
    Integer featureIdx = featureHashes.size();
    if (featureHashes.containsKey(fHash)) {
      featureIdx = featureHashes.get(fHash);
    } else {
      featureHashes.put(fHash, featureIdx);
    }
    return featureIdx;
  }
  
  /**
   * This adds a feature with a given count to an instance. If the feature 
   * already exists, the counts are added.
   * 
   * @param instance the instance this feature is added to.
   * @param fStr the feature.
   * @param fCount the count of the feature.
   */
  private void addFeature(Instance.Builder instance, String fStr, 
        double fCount) {
    SortedMap<Integer, Double> feat2val = new TreeMap<Integer, Double>();
    for (Feature f : instance.getFeatureList()) {
      feat2val.put(f.getIndex(), f.getValue());
    }
    int featureIdx;
    if (useHash) {
      featureIdx = getFeatureHashIdx(fStr);
    } else {
      featureIdx = getFeatureIdx(fStr);
    }
    if (feat2val.containsKey(featureIdx)) {
      // Add up feature count.
      fCount += feat2val.get(featureIdx);
    }
    feat2val.put(featureIdx, fCount);
    List<Feature> featureList = new ArrayList<Feature>(feat2val.size());
    for (Entry<Integer, Double> fe : feat2val.entrySet()) {
      featureList.add(
          Feature.newBuilder().setIndex(fe.getKey()).setValue(fe.getValue())
          .build());
    }
    instance.clearFeature();
    instance.addAllFeature(featureList);
  }

  /**
   * This adds features to an instance, each with a count of 1.0. If a feauture
   * exists already, counts are added.
   * 
   * @param instance the instance features are added to.
   * @param featureStrings the features.
   */
  private void addFeatures(Instance.Builder instance, String[] featureStrings) {
    SortedMap<Integer, Double> feat2val = new TreeMap<Integer, Double>();
    for (Feature f : instance.getFeatureList()) {
      feat2val.put(f.getIndex(), f.getValue());
    }
    for (String fStr : featureStrings) {
      Double fCount = 1.0;
      int featureIdx;
      if (useHash) {
        featureIdx = getFeatureHashIdx(fStr);
      } else {
        featureIdx = getFeatureIdx(fStr);
      }
      if (feat2val.containsKey(featureIdx)) {
        // Add up feature count.
        fCount += feat2val.get(featureIdx);
      }
      feat2val.put(featureIdx, fCount);
    }
    List<Feature> featureList = new ArrayList<Feature>(feat2val.size());
    for (Entry<Integer, Double> fe : feat2val.entrySet()) {
      featureList.add(
          Feature.newBuilder().setIndex(fe.getKey()).setValue(fe.getValue())
          .build());
    }
    instance.clearFeature();
    instance.addAllFeature(featureList);
  }
  
  /**
   * This extracts instances from a document. One instance is returned for every
   * annotated matching relation.
   * 
   * The following features are extracted:
   * - every token as a feature
   * - the shortest path in the dependency tree as a feature
   * - context around the matching argument ARG1 and ARG2, specifically:
   *   - ARG1 + middle_tokens + ARG2
   *   - one_token + ARG1
   *   - two_tokens + ARG1
   *   - ARG2 + one_token
   *   - ARG2 + two_tokens
   * 
   * The arguments are wild-carded.
   * 
   * @param doc the document to extract instances from.
   * @return the training instances with features and the matching relation as
   * a label.
   */
  // TODO: refactor
  public List<Instance> extract(Document doc) {
    List<Instance> instances = new ArrayList<Instance>();    
    // Extract for all binary properties.
    for (CompoundGroup cg : doc.getCompoundList()) {
      if (cg.getType() == AnnotationType.PROPERTY) {
        for (Compound c : cg.getCompoundList()) {
          if (c.getSlotCount() == 2) {
            int argIdx1 = c.getSlot(0).getTokenIndex();
            int argIdx2 = c.getSlot(1).getTokenIndex();
            int leftIdx = Math.min(argIdx1, argIdx2);
            int rightIdx = Math.max(argIdx1, argIdx2);
            StringBuffer sb = new StringBuffer();
            for (int i = Math.max(0, leftIdx - 2); i < leftIdx; ++i) {
              sb.append(doc.getToken(i).getText());
              if (i + 1 < leftIdx) {
                sb.append(' ');
              }
            }
            String leftTwoTokens = sb.toString();
            sb = new StringBuffer();
            for (int i = Math.max(0, leftIdx - 1); i < leftIdx; ++i) {
              sb.append(doc.getToken(i).getText());
              if (i + 1 < leftIdx) {
                sb.append(' ');
              }
            }
            String leftToken = sb.toString();
            sb = new StringBuffer();
            for (int i = rightIdx + 1; i < Math.min(rightIdx + 3, 
                doc.getTokenCount()); ++i) {
              if (i > rightIdx + 1) {
                sb.append(' ');
              }
              sb.append(doc.getToken(i).getText());
            }
            String rightTwoTokens = sb.toString();
            sb = new StringBuffer();
            for (int i = rightIdx + 1; i < Math.min(rightIdx + 2, 
                doc.getTokenCount()); ++i) {
              if (i > rightIdx + 1) {
                sb.append(' ');
              }
              sb.append(doc.getToken(i).getText());
            }
            String rightToken = sb.toString();
            sb = new StringBuffer();
            for (int i = leftIdx + 1; i < rightIdx; ++i) {
              if (i > leftIdx + 1) {
                sb.append(' ');
              }
              sb.append(doc.getToken(i).getText());
            }
            String middleTokens = sb.toString();
            Instance.Builder instance = Instance.newBuilder();
            String leftArg = argIdx1 < argIdx2 ? "ARG1" : "ARG2";
            String rightArg = argIdx1 > argIdx2 ? "ARG1" : "ARG2";

            String[] featureStrings = new String[]{
                leftArg + " " + middleTokens + " " + rightArg,
                leftToken + " " + leftArg,
                rightArg + " " + rightToken,
                leftTwoTokens + " " + leftArg,
                rightArg + " " + rightTwoTokens
            };
            addFeatures(instance, featureStrings);

            String[] tokenFeatures = new String[doc.getTokenCount() - 2];
            for (int i=0, featureIndex=0; i < doc.getTokenCount(); ++i) {
              if (i == argIdx1 || i == argIdx2) {
                continue;
              }
              tokenFeatures[featureIndex] = doc.getToken(i).getText();
              ++featureIndex;
            }
            addFeatures(instance, tokenFeatures);
            
            DependencyTree t = new DependencyTree(doc);
            addFeature(instance, t.pathString(argIdx1, argIdx2), 1.0);

            setLabel(instance, c.getText());
            instances.add(instance.build());
          }
        }
      }
    }
    return instances;
  }

  /**
   * This returns the feature map, that maps features to indices.
   * 
   * @return the feature map.
   */
  public FeatureMap getFeatureMap() {
    FeatureMap.Builder fm = FeatureMap.newBuilder();
    if (useHash) {
      BiMap<Integer, Integer> idx2feat = featureHashes.inverse();
      for (int i = 0; i < featureHashes.size(); ++i) {
        fm.addFeatureHash(idx2feat.get(i));
      }      
    } else {
      BiMap<Integer, String> idx2feat = features.inverse();
      for (int i = 0; i < features.size(); ++i) {
        fm.addFeature(idx2feat.get(i));
      }
    }
    BiMap<Integer, String> idx2label = labels.inverse();
    for (int i = 0; i < labels.size(); ++i) {
      fm.addCategory(idx2label.get(i));
    }
    return fm.build();
  }
}
